import os

import tensorflow as tf
from keras.datasets import mnist
from tensorflow.python.keras import layers, optimizers, Sequential, losses

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

# 加载数据集 60k,10k
(x_train, y_train), (x_test, y_test) = mnist.load_data()
print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)

# one-hot
y_train = tf.one_hot(y_train, depth=10)
y_test = tf.one_hot(y_test, depth=10)

# cnn 主要是处理图片数据，输入格式：b,h,w,c
# (60000, 28, 28）------>（60000，28， 28， 1）
# x_train = tf.reshape(x_train, (x_train.shape[0], x_train.shape[1], x_train.shape[2], 1))
# x_test = tf.reshape(x_test, (x_test.shape[0], x_test.shape[1], x_test.shape[2], 1))
# 还可以用
x_train = tf.expand_dims(x_train, axis=-1)
x_test = tf.expand_dims(x_test, axis=-1)
print(x_train.shape)
print(x_test.shape)

# 构建网络结构 （新方式）
network = Sequential([
    # 卷积层 6个卷积核 卷积核大小3*3  s:1   pad=same
    layers.Conv2D(filters=6, kernel_size=3, padding='same', strides=1, activation='sigmoid', input_shape=(28, 28, 1)),

    # 池化层 s:2  池化核2*2
    layers.MaxPool2D(pool_size=2, strides=2),

    # 卷积层 16个卷积核  卷积核大小5*5
    layers.Conv2D(filters=16, kernel_size=3, activation='sigmoid'),

    # 池化层
    layers.MaxPool2D(pool_size=2, strides=2),
    # layers.Conv2D(filters=120, kernel_size=5, activation='sigmoid'),

    # 数据打平，相当于转换成一个长向量，方便全连接层处理
    layers.Flatten(),

    # 全连接层
    layers.Dense(120, activation='sigmoid'),
    layers.Dense(84, activation='sigmoid'),
    layers.Dense(10, activation='sigmoid'),
])

network.summary()

# 加载网络以及训练
network.compile(
    optimizer=optimizers.gradient_descent_v2.SGD(learning_rate=0.6),
    loss=losses.categorical_crossentropy,
    metrics=['accuracy']
)
# 统计完成率
network.fit(x_train, y_train, epochs=30, batch_size=128, validation_split=0.1)

acc = network.evaluate(x_test, y_test, verbose=1)
print(f'test acc:{acc[1]}')
